import time
import cv2
import numpy as np
import math
import pytesseract
from PySide6.QtCore import QObject, QThread, Signal


# @brief: 预处理，统计表格中横线的倾斜角给图像进行姿态矫正
# @param:
#   img_color	输入的彩色图像
# @retval:
#   img_copy	输出矫正后的图像
def preProcess(img_color):
    img_copy = img_color.copy()
    img_gray = cv2.cvtColor(img_color, cv2.COLOR_BGR2GRAY)  # 转灰度图
    edges = cv2.Canny(img_gray, 50, 150, apertureSize=3)  # 检测边缘
    # cv2.imwrite("./results/edges.jpg", edges)
    # 检测直线 (输入图像 搜索半径步进 角度步进 累计点阈值 最短线段长度 最短连线长度)
    lines = cv2.HoughLinesP(edges,
                            1,
                            np.pi / 180,
                            500,
                            0,
                            minLineLength=500,
                            maxLineGap=100)

    # 计算直线的平均角度
    theta_total = 0
    theta_count = 0
    for line in lines:
        x1, y1, x2, y2 = line[0]
        # rho = np.sqrt((x1 - x2)**2 + (y1 - y2)**2)
        theta = math.atan(float(y2 - y1) / float(x2 - x1 + 0.001))
        # print(rho, theta, x1, y1, x2, y2)
        if theta < np.pi / 4 and theta > -np.pi / 4:  # 筛选 -90° 到 90° 的直线
            theta_total = theta_total + theta
            theta_count += 1
            # cv2.line(img_copy, (x1, y1), (x2, y2), (0, 0, 255), 2)		# 将检测到的直线画到原图上
    theta_average = theta_total / theta_count  # 直线的平均角度
    # radius_average = theta_average * 180 / np.pi
    # print(f"theta average angle:{theta_average}, radius average rad:{radius_average}")
    # cv2.imwrite('./results/hough_lines.jpg', img_copy)

    # 仿射变换矫正图像
    rows, cols, channels = img_color.shape
    affineShrinkTranslationRotation = cv2.getRotationMatrix2D(
        (0, rows), theta_average * 180 / np.pi, 1)
    ShrinkTranslationRotation = cv2.warpAffine(
        img_copy, affineShrinkTranslationRotation, (cols, rows))
    img_copy = cv2.warpAffine(img_copy, affineShrinkTranslationRotation,
                              (cols, rows))
    # cv2.imwrite('./results/image_Rotation.jpg', ShrinkTranslationRotation)
    return img_copy


# @brief: 分割表格中的矩形
# @param:
#   img_color	输入的彩色图像
#   minArea     表格矩形框的最小面积
#   minHeight   表格最小高度
#   maxHeight   表格最大高度
#   minWidth    表格最小宽度
# @retval:
#   rects       返回表格矩形框
#   Intersections   返回矩形框顶点
def spiltRect(img_color, minArea, minHeight, maxHeight, minWidth):
    img_copy = img_color.copy()
    img_gray = cv2.cvtColor(img_copy, cv2.COLOR_RGB2GRAY)  # 转灰度图
    # img_gray = cv2.GaussianBlur(img_gray, (5,5), 3)

    # 自适应阈值处理转为二值图
    binary_img = cv2.adaptiveThreshold(~img_gray, 255,
                                       cv2.ADAPTIVE_THRESH_MEAN_C,
                                       cv2.THRESH_BINARY, 15, -2)  # 注意 ~ 符号
    # cv2.imwrite('./results/binary_img.jpg', binary_img)

    # 腐蚀操作，长横条内核处理
    rows, cols = img_gray.shape
    scale = 44  # 20, 50, 45, 40, 44
    horizontalsize = cols / scale
    horizontalStructure = cv2.getStructuringElement(
        cv2.MORPH_RECT, (int(horizontalsize), 1))  # 这里要做个类型转换，切片必须是整数
    erosion_horizon = cv2.erode(binary_img, horizontalStructure, iterations=1)
    dilation_horizon = cv2.dilate(erosion_horizon,
                                  horizontalStructure,
                                  iterations=1)

    # 腐蚀操作，长竖条内核处理
    scale = 39
    # 20, 50, 45, 40, 39
    verticalsize = rows / scale
    verticalStructure = cv2.getStructuringElement(cv2.MORPH_RECT,
                                                  (1, int(verticalsize)))
    erosion_vertical = cv2.erode(binary_img, verticalStructure, iterations=1)
    dilation_vertical = cv2.dilate(erosion_vertical,
                                   verticalStructure,
                                   iterations=1)

    # 横线图与竖线图叠加
    mask = dilation_horizon + dilation_vertical
    # cv2.imwrite('./results/mask.jpg', mask)
    joints = cv2.bitwise_and(dilation_horizon,
                             dilation_vertical)  # 与操作，表格的交点处会留下几个紧凑的像素点
    # cv2.imwrite('./results/joints.jpg', joints)

    # 形态学闭运算，先膨胀后腐蚀，将表格交点处的几个点连起来
    k1 = np.ones((5, 5), np.uint8)  # 指定核大小
    Intersections = cv2.morphologyEx(joints, cv2.MORPH_CLOSE, k1)
    # cv2.imwrite('./results/Intersections.jpg', Intersections)

    # 查找图像矩形轮廓
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_LIST,
                                           cv2.CHAIN_APPROX_SIMPLE)
    contours_num = len(contours)
    print(f"contours length:{contours_num}")
    # cv2.drawContours(img_copy, contours, -1, (0, 0, 255), 3)
    # cv2.imwrite('./results/coutours.jpg', img_copy)

    # 根据矩形大小筛选矩形框
    rects = []  # 每个表格的矩形框
    for i in range(contours_num):
        contour = contours[i]
        area = cv2.contourArea(contour)
        if area < minArea:
            continue
        # 根据四个顶点近似矩形
        approx = cv2.approxPolyDP(contour, 3, True)  #3
        x, y, w, h = cv2.boundingRect(approx)
        rect = (x, y, w, h)
        #cv.rectangle(img, (x, y), (x+w, y+h), (0, 0, 255), 3)
        roi = Intersections[y:y + h, x:x + w]
        Intersections_contours, Intersections_hierarchy = cv2.findContours(
            roi, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        #print len(joints_contours)

        # 根据大小筛选矩形框
        if h < maxHeight and h > minHeight and w > minWidth and len(
                Intersections_contours) <= 4:  # important
            cv2.rectangle(img_copy, (x, y), (x + w, y + h), (0, 0, 255), 3)
            rects.append(rect)
    # cv2.imwrite('./results/table_out.jpg', img_copy)
    print(f"rects number:{len(rects)}")
    return rects, Intersections


# @brief: 在分割的矩形区域截图
# @param:
#   img_color	输入的彩色图像
#   rects       输入表格矩形框
#   rho         行与行之间的最小间隔
# @retval:
#   imgs        返回分割好的图像
def getROI(img_color, rects, rho=10):
    imgs = []

    rects = sorted(rects, key=lambda x: x[1], reverse=False)  # y 方向排序，升序

    # y 方向分组，y 相近的认为是同一行表格
    sort_rects = []  # 排序好的表格
    rect_array = []  # 表格的所有行
    rect_line = []  # 每一行表格的矩形框
    for i in range(len(rects) - 1):
        # 处理第一次和最后一次
        if i == 0:
            rect_line.append(rects[0])
        elif i == len(rects) - 2:
            rect_array.append(rect_line)

        # 行间距小的认为是同一行
        rect1 = rects[i]
        rect2 = rects[i + 1]
        if (rect2[1] - rect1[1] < rho):
            rect_line.append(rect2)
        else:
            rect_array.append(rect_line)
            rect_line = []  # clear
            rect_line.append(rect2)

    # 对每一行表格按 x 方向排序，并将结果保存到 sort_rects
    for line in rect_array:
        sort_line = sorted(line, key=lambda x: x[0], reverse=False)  # x 升序
        for rect in sort_line:
            sort_rects.append(rect)

    count = 0
    # 根据矩形框截图
    for rect in sort_rects:
        count += 1
        x = rect[0]
        y = rect[1]
        w = rect[2]
        h = rect[3]
        img = img_color[y + 20:y + h - 20,
                        x + 20:x + w - 20].copy()  # 每个边界往里缩小一圈，把表格框线去掉
        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 转为灰度图
        __, img_gray = cv2.threshold(img_gray, 125, 255,
                                     cv2.THRESH_BINARY)  # 转为二值图
        img_gray = cv2.copyMakeBorder(img_gray,
                                      200,
                                      200,
                                      200,
                                      200,
                                      cv2.BORDER_CONSTANT,
                                      value=255)  # 扩展图像边界，增加 tesseract 识别率
        imgs.append(img_gray)
        # filename = f"./results/imgs/img{count}.jpg"
        # cv2.imwrite(filename, img_gray)
    return imgs


# @brief: 分割表格，每一个单元格为一张图片
# @param:
#   filename	表格图片路径
#   minArea     表格矩形框的最小面积
#   minHeight   表格最小高度
#   maxHeight   表格最大高度
#   minWidth    表格最小宽度
# @retval:
#   imgs        返回分割好的图像
#   rects       返回表格矩形框
#   Intersections   返回矩形框顶点
def table_to_grid(filename,
                  minArea=500,
                  minHeight=50,
                  maxHeight=300,
                  minWidth=50):
    "@brief: 分割表格，每一个单元格为一张图片\n@param: \n\nfilename	表格图片路径 \n\nminArea     表格矩形框的最小面积 \n\nminHeight   表格最小高度 \n\nmaxHeight   表格最大高度 \n\nminWidth    表格最小宽度 \n@retval: \n\nimgs        分割好的图像 \n\nrects       表格矩形框 \n\nIntersections   矩形框顶点"
    image = cv2.imread(filename)
    image1 = preProcess(image)
    rects, Intersections = spiltRect(image1, minArea, minHeight, maxHeight,
                                     minWidth)
    imgs = getROI(image1, rects, minHeight)
    return imgs, rects, Intersections


# 表格识别对象
class GridWorker(QObject):

    # slots
    resultReady = Signal(str)

    def __init__(self, minArea=500, minHeight=50, maxHeight=300, minWidth=50):
        super().__init__()
        self.minArea = minArea
        self.minHeight = minHeight
        self.maxHeight = maxHeight
        self.minWidth = minWidth

    def openImage(self, image_path):
        self.image_path = image_path

    def table2str(self):
        imgs, rects, Intersections = table_to_grid(self.image_path, 3000, 50,
                                                   500, 50)
        for img in imgs:
            info = pytesseract.image_to_string(img, lang='chi_sim', timeout=5)
            self.resultReady.emit(info)


if __name__ == '__main__':
    image_path = "D:\\MyProject\\Python Project\\grid_spilt\\samples\\grid_src.jpg"
    imgs, rects, Intersections = table_to_grid(image_path, 3000, 50, 500, 50)
    count = 0
    for img in imgs:
        count += 1
        # img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # ret, img_gray = cv2.threshold(img_gray, 125, 255, cv2.THRESH_BINARY)
        # cv2.imwrite(filename, img_gray)
        # info = pytesseract.image_to_string(img_gray, lang = 'chi_sim', timeout=5)
        info = pytesseract.image_to_string(img, lang='chi_sim', timeout=5)
        print(info)
